import mediapipe as mp
import cv2
import numpy as np


class MediaPipeWrapper:
    def __init__(self, min_detection_confidence=0.5, model_complexity=1):
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            min_detection_confidence=min_detection_confidence,
            model_complexity=model_complexity,
        )
        self.mp_drawing = mp.solutions.drawing_utils
        self.drawing_styles = mp.solutions.drawing_styles

        self.keypoint_colors = {
            "head": (255, 0, 0),  # 蓝色
            "torso": (0, 255, 0),  # 绿色
            "arms": (255, 165, 0),  # 橙色
            "legs": (255, 0, 255),  # 紫色
        }

    def process_frame(self, frame):
        # 将 BGR 图像转换为 RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image.flags.writeable = False  # 为了提高性能，将图像标记为只读
        results = self.pose.process(image)

        # 将图像标记为可写，并转换回 BGR
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        # 绘制姿势检测结果
        if results.pose_landmarks:
            # 增强可视化效果
            image = self.draw_enhanced_visualization(image, results.pose_landmarks)

            # 添加姿态轨迹
            image = self.draw_motion_trajectory(image, results.pose_landmarks)

        return image, results

    def draw_enhanced_visualization(self, image, landmarks):
        """绘制增强的可视化效果"""
        h, w = image.shape[:2]

        # 绘制骨架连接
        connections = [
            # 躯干
            ([23, 24], self.keypoint_colors["torso"]),  # 髋部
            ([11, 12], self.keypoint_colors["torso"]),  # 肩部
            ([11, 23], self.keypoint_colors["torso"]),  # 左侧躯干
            ([12, 24], self.keypoint_colors["torso"]),  # 右侧躯干
            # 手臂
            ([11, 13, 15], self.keypoint_colors["arms"]),  # 左臂
            ([12, 14, 16], self.keypoint_colors["arms"]),  # 右臂
            # 腿部
            ([23, 25, 27], self.keypoint_colors["legs"]),  # 左腿
            ([24, 26, 28], self.keypoint_colors["legs"]),  # 右腿
        ]

        # 绘制连接线和关节点
        for connection, color in connections:
            for i in range(len(connection) - 1):
                start_idx = connection[i]
                end_idx = connection[i + 1]

                start_point = (
                    int(landmarks.landmark[start_idx].x * w),
                    int(landmarks.landmark[start_idx].y * h),
                )
                end_point = (
                    int(landmarks.landmark[end_idx].x * w),
                    int(landmarks.landmark[end_idx].y * h),
                )

                cv2.line(image, start_point, end_point, color, 2)
                cv2.circle(image, start_point, 4, color, -1)
                cv2.circle(image, end_point, 4, color, -1)

        return image

    def draw_motion_trajectory(self, image, landmarks):
        """绘制运动轨迹"""
        if not hasattr(self, "trajectory_points"):
            self.trajectory_points = []

        # 获取关键点位置（以鼻尖为例）
        h, w = image.shape[:2]
        nose_point = (
            int(landmarks.landmark[0].x * w),
            int(landmarks.landmark[0].y * h),
        )

        # 更新轨迹点
        self.trajectory_points.append(nose_point)
        if len(self.trajectory_points) > 30:  # 保持最近30帧的轨迹
            self.trajectory_points.pop(0)

        # 绘制轨迹
        for i in range(1, len(self.trajectory_points)):
            cv2.line(
                image,
                self.trajectory_points[i - 1],
                self.trajectory_points[i],
                (0, 255, 255),  # 黄色
                2,
            )

        return image
